// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <ktest/ktest.h>

#include <string>

#include <nebula/core/array.h>
#include <nebula/array/builder_nested.h>
#include <nebula/array/builder_union.h>
// TODO ipc shouldn't be included here
#include <tests/ipc/test_common.h>
#include <nebula/testing/builder.h>
#include <nebula/testing/ktest_util.h>
#include <nebula/testing/util.h>
#include <nebula/types/type.h>
#include <turbo/base/checked_cast.h>

namespace nebula {


    TEST(TestUnionArray, TestSliceEquals) {
        std::shared_ptr<RecordBatch> batch;
        ASSERT_OK(ipc::test::MakeUnion(&batch));

        auto CheckUnion = [](std::shared_ptr<Array> array) {
            const int64_t size = array->length();
            std::shared_ptr<Array> slice, slice2;
            slice = array->slice(2);
            ASSERT_EQ(size - 2, slice->length());

            slice2 = array->slice(2);
            ASSERT_EQ(size - 2, slice->length());

            ASSERT_TRUE(slice->equals(slice2));
            ASSERT_TRUE(array->range_equals(2, array->length(), 0, slice));

            // Chained slices
            slice2 = array->slice(1)->slice(1);
            ASSERT_TRUE(slice->equals(slice2));

            slice = array->slice(1, 5);
            slice2 = array->slice(1, 5);
            ASSERT_EQ(5, slice->length());

            ASSERT_TRUE(slice->equals(slice2));
            ASSERT_TRUE(array->range_equals(1, 6, 0, slice));

            AssertZeroPadded(*array);
            TestInitialized(*array);
        };

        CheckUnion(batch->column(0));
        CheckUnion(batch->column(1));
    }

    TEST(TestSparseUnionArray, GetFlattenedField) {
        auto ty = sparse_union({field("ints", int64()), field("strs", utf8())}, {2, 7});
        auto ints = assert_array_from_json(int64(), "[0, 1, 2, 3]");
        auto strs = assert_array_from_json(utf8(), R"(["a", null, "c", "d"])");
        auto ids = assert_array_from_json(int8(), "[2, 7, 2, 7]")->data()->buffers[1];
        const int length = 4;

        {
            SparseUnionArray arr(ty, length, {ints, strs}, ids);
            ASSERT_OK(arr.validate_full());

            ASSERT_OK_AND_ASSIGN(auto flattened, arr.GetFlattenedField(0));
            AssertArraysEqual(*assert_array_from_json(int64(), "[0, null, 2, null]"), *flattened,
                    /*verbose=*/true);

            ASSERT_OK_AND_ASSIGN(flattened, arr.GetFlattenedField(1));
            AssertArraysEqual(*assert_array_from_json(utf8(), R"([null, null, null, "d"])"), *flattened,
                    /*verbose=*/true);

            const auto sliced = turbo::checked_pointer_cast<SparseUnionArray>(arr.slice(1, 2));

            ASSERT_OK_AND_ASSIGN(flattened, sliced->GetFlattenedField(0));
            AssertArraysEqual(*assert_array_from_json(int64(), "[null, 2]"), *flattened, /*verbose=*/true);

            ASSERT_OK_AND_ASSIGN(flattened, sliced->GetFlattenedField(1));
            AssertArraysEqual(*assert_array_from_json(utf8(), R"([null, null])"), *flattened,
                    /*verbose=*/true);

            ASSERT_RAISES(turbo::StatusCode::kInvalidArgument, arr.GetFlattenedField(-1));
            ASSERT_RAISES(turbo::StatusCode::kInvalidArgument, arr.GetFlattenedField(2));
        }
        {
            SparseUnionArray arr(ty, length - 2, {ints->slice(1, 2), strs->slice(1, 2)}, ids);
            ASSERT_OK(arr.validate_full());

            ASSERT_OK_AND_ASSIGN(auto flattened, arr.GetFlattenedField(0));
            AssertArraysEqual(*assert_array_from_json(int64(), "[1, null]"), *flattened, /*verbose=*/true);

            ASSERT_OK_AND_ASSIGN(flattened, arr.GetFlattenedField(1));
            AssertArraysEqual(*assert_array_from_json(utf8(), R"([null, "c"])"), *flattened,
                    /*verbose=*/true);

            const auto sliced = turbo::checked_pointer_cast<SparseUnionArray>(arr.slice(1, 1));

            ASSERT_OK_AND_ASSIGN(flattened, sliced->GetFlattenedField(0));
            AssertArraysEqual(*assert_array_from_json(int64(), "[null]"), *flattened, /*verbose=*/true);

            ASSERT_OK_AND_ASSIGN(flattened, sliced->GetFlattenedField(1));
            AssertArraysEqual(*assert_array_from_json(utf8(), R"(["c"])"), *flattened, /*verbose=*/true);
        }
        {
            SparseUnionArray arr(ty, /*length=*/0, {ints->slice(length), strs->slice(length)},
                                 ids);
            ASSERT_OK(arr.validate_full());

            ASSERT_OK_AND_ASSIGN(auto flattened, arr.GetFlattenedField(0));
            AssertArraysEqual(*assert_array_from_json(int64(), "[]"), *flattened, /*verbose=*/true);

            ASSERT_OK_AND_ASSIGN(flattened, arr.GetFlattenedField(1));
            AssertArraysEqual(*assert_array_from_json(utf8(), "[]"), *flattened,
                    /*verbose=*/true);
        }
    }

    TEST(TestSparseUnionArray, Validate) {
        auto a = assert_array_from_json(int32(), "[4, 5]");
        auto type = sparse_union({field("a", int32())});
        auto children = std::vector<std::shared_ptr<Array>>{a};
        auto type_ids_array = assert_array_from_json(int8(), "[0, 0, 0]");
        auto type_ids = type_ids_array->data()->buffers[1];

        auto arr = std::make_shared<SparseUnionArray>(type, 2, children, type_ids);
        ASSERT_OK(arr->validate_full());
        arr = std::make_shared<SparseUnionArray>(type, 1, children, type_ids,
                /*offset=*/1);
        ASSERT_OK(arr->validate_full());
        arr = std::make_shared<SparseUnionArray>(type, 0, children, type_ids,
                /*offset=*/2);
        ASSERT_OK(arr->validate_full());

        // Length + offset < child length, but it's ok
        arr = std::make_shared<SparseUnionArray>(type, 1, children, type_ids,
                /*offset=*/0);
        ASSERT_OK(arr->validate_full());

        // Length + offset > child length
        arr = std::make_shared<SparseUnionArray>(type, 1, children, type_ids,
                /*offset=*/2);
        ASSERT_RAISES(turbo::StatusCode::kInvalidArgument, arr->validate_full());

        // Offset > child length
        arr = std::make_shared<SparseUnionArray>(type, 0, children, type_ids,
                /*offset=*/3);
        ASSERT_RAISES(turbo::StatusCode::kInvalidArgument, arr->validate_full());
    }

    // -------------------------------------------------------------------------
    // Tests for MakeDense and MakeSparse

    class TestUnionArrayFactories : public ::testing::Test {
    public:
        void SetUp() {
            pool_ = default_memory_pool();
            type_codes_ = {1, 2, 4, 127};
            ArrayFromVector<Int8Type>({0, 1, 2, 0, 1, 3, 2, 0, 2, 1}, &type_ids_);
            ArrayFromVector<Int8Type>({1, 2, 4, 1, 2, 127, 4, 1, 4, 2}, &logical_type_ids_);
            ArrayFromVector<Int8Type>({1, 2, 4, 1, -2, 127, 4, 1, 4, 2}, &invalid_type_ids1_);
            ArrayFromVector<Int8Type>({1, 2, 4, 1, 3, 127, 4, 1, 4, 2}, &invalid_type_ids2_);
        }

        void CheckUnionArray(const UnionArray &array, UnionMode::type mode,
                             const std::vector<std::string> &field_names,
                             const std::vector<int8_t> &type_codes) {
            ASSERT_EQ(mode, array.mode());
            CheckFieldNames(array, field_names);
            CheckTypeCodes(array, type_codes);
            const auto &type_ids = turbo::checked_cast<const Int8Array &>(*type_ids_);
            for (int64_t i = 0; i < type_ids.length(); ++i) {
                ASSERT_EQ(array.child_id(i), type_ids.value(i));
            }
            ASSERT_EQ(nullptr, array.field(-1));
            ASSERT_EQ(nullptr, array.field(static_cast<int>(type_ids.length())));
        }

        void CheckFieldNames(const UnionArray &array, const std::vector<std::string> &names) {
            const auto &type = turbo::checked_cast<const UnionType &>(*array.type());
            ASSERT_EQ(type.num_fields(), names.size());
            for (int i = 0; i < type.num_fields(); ++i) {
                ASSERT_EQ(type.field(i)->name(), names[i]);
            }
        }

        void CheckTypeCodes(const UnionArray &array, const std::vector<int8_t> &codes) {
            const auto &type = turbo::checked_cast<const UnionType &>(*array.type());
            ASSERT_EQ(codes, type.type_codes());
        }

    protected:
        MemoryPool *pool_;
        std::vector<int8_t> type_codes_;
        std::shared_ptr<Array> type_ids_;
        std::shared_ptr<Array> logical_type_ids_;
        std::shared_ptr<Array> invalid_type_ids1_;
        std::shared_ptr<Array> invalid_type_ids2_;
    };

    TEST_F(TestUnionArrayFactories, TestMakeDense) {
        std::shared_ptr<Array> value_offsets;
        // type_ids_:                       {0, 1, 2, 0, 1, 3, 2, 0, 2, 1}
        ArrayFromVector<Int32Type, int32_t>({0, 0, 0, 1, 1, 0, 1, 2, 1, 2}, &value_offsets);

        auto children = std::vector<std::shared_ptr<Array>>(4);
        ArrayFromVector<StringType, std::string>({"abc", "def", "xyz"}, &children[0]);
        ArrayFromVector<UInt8Type>({10, 20, 30}, &children[1]);
        ArrayFromVector<Fp64Type>({1.618, 2.718, 3.142}, &children[2]);
        ArrayFromVector<Int8Type>({-12}, &children[3]);

        std::vector<std::string> field_names = {"str", "int1", "real", "int2"};

        std::shared_ptr<Array> result;
        const UnionArray *union_array;

        // without field names and type codes
        ASSERT_OK_AND_ASSIGN(result,
                             DenseUnionArray::create(*type_ids_, *value_offsets, children));
        ASSERT_OK(result->validate_full());
        union_array = turbo::checked_cast<const UnionArray *>(result.get());
        CheckUnionArray(*union_array, UnionMode::DENSE, {"0", "1", "2", "3"}, {0, 1, 2, 3});

        // with field name
        ASSERT_RAISES(turbo::StatusCode::kInvalidArgument,
                      DenseUnionArray::create(*type_ids_, *value_offsets, children, {"one"}));
        ASSERT_OK_AND_ASSIGN(
                result, DenseUnionArray::create(*type_ids_, *value_offsets, children, field_names));
        ASSERT_OK(result->validate_full());
        union_array = turbo::checked_cast<const UnionArray *>(result.get());
        CheckUnionArray(*union_array, UnionMode::DENSE, field_names, {0, 1, 2, 3});

        // with type codes
        ASSERT_RAISES(turbo::StatusCode::kInvalidArgument, DenseUnionArray::create(*logical_type_ids_, *value_offsets,
                                                                                   children, std::vector<int8_t>{0}));
        ASSERT_OK_AND_ASSIGN(result, DenseUnionArray::create(*logical_type_ids_, *value_offsets,
                                                             children, type_codes_));
        ASSERT_OK(result->validate_full());
        union_array = turbo::checked_cast<const UnionArray *>(result.get());
        CheckUnionArray(*union_array, UnionMode::DENSE, {"0", "1", "2", "3"}, type_codes_);

        // with field names and type codes
        ASSERT_RAISES(turbo::StatusCode::kInvalidArgument, DenseUnionArray::create(*logical_type_ids_, *value_offsets,
                                                                                   children, {"one"}, type_codes_));
        ASSERT_OK_AND_ASSIGN(result, DenseUnionArray::create(*logical_type_ids_, *value_offsets,
                                                             children, field_names, type_codes_));
        ASSERT_OK(result->validate_full());
        union_array = turbo::checked_cast<const UnionArray *>(result.get());
        CheckUnionArray(*union_array, UnionMode::DENSE, field_names, type_codes_);

        // turbo::StatusCode::kInvalidArgument type codes
        ASSERT_OK_AND_ASSIGN(result, DenseUnionArray::create(*invalid_type_ids1_, *value_offsets,
                                                             children, type_codes_));
        ASSERT_RAISES(turbo::StatusCode::kInvalidArgument, result->validate_full());
        ASSERT_OK_AND_ASSIGN(result, DenseUnionArray::create(*invalid_type_ids2_, *value_offsets,
                                                             children, type_codes_));
        ASSERT_RAISES(turbo::StatusCode::kInvalidArgument, result->validate_full());

        // turbo::StatusCode::kInvalidArgument offsets
        // - offset out of bounds at index 5
        std::shared_ptr<Array> invalid_offsets;
        ArrayFromVector<Int32Type, int32_t>({0, 0, 0, 1, 1, 1, 1, 2, 1, 2}, &invalid_offsets);
        ASSERT_OK_AND_ASSIGN(result,
                             DenseUnionArray::create(*type_ids_, *invalid_offsets, children));
        ASSERT_RAISES(turbo::StatusCode::kInvalidArgument, result->validate_full());
        // - negative offset at index 5
        ArrayFromVector<Int32Type, int32_t>({0, 0, 0, 1, 1, -1, 1, 2, 1, 2}, &invalid_offsets);
        ASSERT_OK_AND_ASSIGN(result,
                             DenseUnionArray::create(*type_ids_, *invalid_offsets, children));
        ASSERT_RAISES(turbo::StatusCode::kInvalidArgument, result->validate_full());
        // - non-monotonic offset at index 3
        ArrayFromVector<Int32Type, int32_t>({1, 0, 0, 0, 1, 0, 1, 2, 1, 2}, &invalid_offsets);
        ASSERT_OK_AND_ASSIGN(result,
                             DenseUnionArray::create(*type_ids_, *invalid_offsets, children));
        ASSERT_RAISES(turbo::StatusCode::kInvalidArgument, result->validate_full());
    }

    TEST_F(TestUnionArrayFactories, TestMakeSparse) {
        auto children = std::vector<std::shared_ptr<Array>>(4);
        ArrayFromVector<StringType, std::string>(
                {"abc", "", "", "def", "", "", "", "xyz", "", ""}, &children[0]);
        ArrayFromVector<UInt8Type>({0, 10, 0, 0, 20, 0, 0, 0, 0, 30}, &children[1]);
        ArrayFromVector<Fp64Type>({0.0, 0.0, 1.618, 0.0, 0.0, 0.0, 2.718, 0.0, 3.142, 0.0},
                                    &children[2]);
        ArrayFromVector<Int8Type>({0, 0, 0, 0, 0, -12, 0, 0, 0, 0}, &children[3]);

        std::vector<std::string> field_names = {"str", "int1", "real", "int2"};

        std::shared_ptr<Array> result;

        // without field names and type codes
        ASSERT_OK_AND_ASSIGN(result, SparseUnionArray::create(*type_ids_, children));
        ASSERT_OK(result->validate_full());
        CheckUnionArray(turbo::checked_cast<UnionArray &>(*result), UnionMode::SPARSE,
                        {"0", "1", "2", "3"}, {0, 1, 2, 3});

        // with field names
        ASSERT_RAISES(turbo::StatusCode::kInvalidArgument, SparseUnionArray::create(*type_ids_, children, {"one"}));
        ASSERT_OK_AND_ASSIGN(result, SparseUnionArray::create(*type_ids_, children, field_names));
        ASSERT_OK(result->validate_full());
        CheckUnionArray(turbo::checked_cast<UnionArray &>(*result), UnionMode::SPARSE, field_names,
                        {0, 1, 2, 3});

        // with type codes
        ASSERT_RAISES(turbo::StatusCode::kInvalidArgument, SparseUnionArray::create(*logical_type_ids_, children,
                                                                                    std::vector<int8_t>{0}));
        ASSERT_OK_AND_ASSIGN(result,
                             SparseUnionArray::create(*logical_type_ids_, children, type_codes_));
        ASSERT_OK(result->validate_full());
        CheckUnionArray(turbo::checked_cast<UnionArray &>(*result), UnionMode::SPARSE,
                        {"0", "1", "2", "3"}, type_codes_);

        // with field names and type codes
        ASSERT_RAISES(turbo::StatusCode::kInvalidArgument,
                      SparseUnionArray::create(*logical_type_ids_, children, {"one"},
                                               type_codes_));
        ASSERT_OK_AND_ASSIGN(result, SparseUnionArray::create(*logical_type_ids_, children,
                                                              field_names, type_codes_));
        ASSERT_OK(result->validate_full());
        CheckUnionArray(turbo::checked_cast<UnionArray &>(*result), UnionMode::SPARSE, field_names,
                        type_codes_);

        // turbo::StatusCode::kInvalidArgument type codes
        ASSERT_OK_AND_ASSIGN(
                result, SparseUnionArray::create(*invalid_type_ids1_, children, type_codes_));
        ASSERT_RAISES(turbo::StatusCode::kInvalidArgument, result->validate_full());
        ASSERT_OK_AND_ASSIGN(
                result, SparseUnionArray::create(*invalid_type_ids2_, children, type_codes_));
        ASSERT_RAISES(turbo::StatusCode::kInvalidArgument, result->validate_full());

        // turbo::StatusCode::kInvalidArgument child length
        ArrayFromVector<Int8Type>({0, 0, 0, 0, 0, -12, 0, 0, 0}, &children[3]);
        ASSERT_RAISES(turbo::StatusCode::kInvalidArgument, SparseUnionArray::create(*type_ids_, children));
    }

    template<typename B>
    class UnionBuilderTest : public ::testing::Test {
    public:
        int8_t I8 = 8, STR = 13, DBL = 7;

        virtual void AppendInt(int8_t i) {
            expected_types_vector.push_back(I8);
            ASSERT_OK(union_builder->append(I8));
            ASSERT_OK(i8_builder->append(i));
        }

        virtual void AppendString(const std::string &str) {
            expected_types_vector.push_back(STR);
            ASSERT_OK(union_builder->append(STR));
            ASSERT_OK(str_builder->append(str));
        }

        virtual void AppendDouble(double dbl) {
            expected_types_vector.push_back(DBL);
            ASSERT_OK(union_builder->append(DBL));
            ASSERT_OK(dbl_builder->append(dbl));
        }

        void AppendBasics() {
            AppendInt(33);
            AppendString("abc");
            AppendDouble(1.0);
            AppendDouble(-1.0);
            AppendString("");
            AppendInt(10);
            AppendString("def");
            AppendInt(-10);
            AppendDouble(0.5);

            ASSERT_OK(union_builder->finish(&actual));
            ASSERT_OK(actual->validate_full());
            ArrayFromVector<Int8Type, uint8_t>(expected_types_vector, &expected_types);
        }

        void AppendNullsAndEmptyValues() {
            AppendString("abc");
            ASSERT_OK(union_builder->append_null());
            ASSERT_OK(union_builder->append_empty_value());
            expected_types_vector.insert(expected_types_vector.end(), 3, I8);
            AppendInt(42);
            ASSERT_OK(union_builder->append_nulls(2));
            ASSERT_OK(union_builder->append_empty_values(2));
            expected_types_vector.insert(expected_types_vector.end(), 3, I8);

            ASSERT_OK(union_builder->finish(&actual));
            ASSERT_OK(actual->validate_full());
            ArrayFromVector<Int8Type, uint8_t>(expected_types_vector, &expected_types);
        }

        void AppendInferred() {
            I8 = union_builder->AppendChild(i8_builder, "i8");
            ASSERT_EQ(I8, 0);
            AppendInt(33);
            AppendInt(10);

            STR = union_builder->AppendChild(str_builder, "str");
            ASSERT_EQ(STR, 1);
            AppendString("abc");
            AppendString("");
            AppendString("def");
            AppendInt(-10);

            DBL = union_builder->AppendChild(dbl_builder, "dbl");
            ASSERT_EQ(DBL, 2);
            AppendDouble(1.0);
            AppendDouble(-1.0);
            AppendDouble(0.5);

            ASSERT_OK(union_builder->finish(&actual));
            ASSERT_OK(actual->validate_full());
            ArrayFromVector<Int8Type, uint8_t>(expected_types_vector, &expected_types);

            ASSERT_EQ(I8, 0);
            ASSERT_EQ(STR, 1);
            ASSERT_EQ(DBL, 2);
        }

        void AppendListOfInferred(std::shared_ptr<ListArray> *actual) {
            ListBuilder list_builder(default_memory_pool(), union_builder);

            ASSERT_OK(list_builder.append());
            I8 = union_builder->AppendChild(i8_builder, "i8");
            ASSERT_EQ(I8, 0);
            AppendInt(10);

            ASSERT_OK(list_builder.append());
            STR = union_builder->AppendChild(str_builder, "str");
            ASSERT_EQ(STR, 1);
            AppendString("abc");
            AppendInt(-10);

            ASSERT_OK(list_builder.append());
            DBL = union_builder->AppendChild(dbl_builder, "dbl");
            ASSERT_EQ(DBL, 2);
            AppendDouble(0.5);

            ASSERT_OK(list_builder.finish(actual));
            ASSERT_OK((*actual)->validate_full());
            ArrayFromVector<Int8Type, uint8_t>(expected_types_vector, &expected_types);
        }

        std::vector<uint8_t> expected_types_vector;
        std::shared_ptr<Array> expected_types;
        std::shared_ptr<Int8Builder> i8_builder = std::make_shared<Int8Builder>();
        std::shared_ptr<StringBuilder> str_builder = std::make_shared<StringBuilder>();
        std::shared_ptr<Fp64Builder> dbl_builder = std::make_shared<Fp64Builder>();
        std::shared_ptr<B> union_builder = std::make_shared<B>(default_memory_pool());
        std::shared_ptr<UnionArray> actual;
    };

    class DenseUnionBuilderTest : public UnionBuilderTest<DenseUnionBuilder> {
    };

    class SparseUnionBuilderTest : public UnionBuilderTest<SparseUnionBuilder> {
    public:
        using Base = UnionBuilderTest<SparseUnionBuilder>;

        void AppendInt(int8_t i) override {
            Base::AppendInt(i);
            ASSERT_OK(str_builder->append_empty_value());
            ASSERT_OK(dbl_builder->append_empty_value());
        }

        void AppendString(const std::string &str) override {
            Base::AppendString(str);
            ASSERT_OK(i8_builder->append_empty_value());
            ASSERT_OK(dbl_builder->append_empty_value());
        }

        void AppendDouble(double dbl) override {
            Base::AppendDouble(dbl);
            ASSERT_OK(i8_builder->append_empty_value());
            ASSERT_OK(str_builder->append_empty_value());
        }
    };

    TEST_F(DenseUnionBuilderTest, Basics) {
        union_builder.reset(new DenseUnionBuilder(
                default_memory_pool(), {i8_builder, str_builder, dbl_builder},
                dense_union({field("i8", int8()), field("str", utf8()), field("dbl", float64())},
                            {I8, STR, DBL})));
        AppendBasics();

        auto expected_i8 = assert_array_from_json(int8(), "[33, 10, -10]");
        auto expected_str = assert_array_from_json(utf8(), R"(["abc", "", "def"])");
        auto expected_dbl = assert_array_from_json(float64(), "[1.0, -1.0, 0.5]");

        auto expected_offsets = assert_array_from_json(int32(), "[0, 0, 0, 1, 1, 1, 2, 2, 2]");

        ASSERT_OK_AND_ASSIGN(auto expected,
                             DenseUnionArray::create(*expected_types, *expected_offsets,
                                                     {expected_i8, expected_str, expected_dbl},
                                                     {"i8", "str", "dbl"}, {I8, STR, DBL}));

        ASSERT_EQ(expected->type()->to_string(), actual->type()->to_string());
        ASSERT_ARRAYS_EQUAL(*expected, *actual);
    }

    TEST_F(DenseUnionBuilderTest, NullsAndEmptyValues) {
        union_builder.reset(new DenseUnionBuilder(
                default_memory_pool(), {i8_builder, str_builder, dbl_builder},
                dense_union({field("i8", int8()), field("str", utf8()), field("dbl", float64())},
                            {I8, STR, DBL})));
        AppendNullsAndEmptyValues();

        // Four null / empty values (the latter implementation-defined) were appended to I8
        auto expected_i8 = assert_array_from_json(int8(), "[null, 0, 42, null, 0]");
        auto expected_str = assert_array_from_json(utf8(), R"(["abc"])");
        auto expected_dbl = assert_array_from_json(float64(), "[]");

        // "abc", null, 0, 42, null, null, 0, 0
        auto expected_offsets = assert_array_from_json(int32(), "[0, 0, 1, 2, 3, 3, 4, 4]");

        ASSERT_OK_AND_ASSIGN(auto expected,
                             DenseUnionArray::create(*expected_types, *expected_offsets,
                                                     {expected_i8, expected_str, expected_dbl},
                                                     {"i8", "str", "dbl"}, {I8, STR, DBL}));

        ASSERT_EQ(expected->type()->to_string(), actual->type()->to_string());
        ASSERT_ARRAYS_EQUAL(*expected, *actual);
        // Physical arrays must be as expected
        ASSERT_ARRAYS_EQUAL(*expected_i8, *actual->field(0));
        ASSERT_ARRAYS_EQUAL(*expected_str, *actual->field(1));
        ASSERT_ARRAYS_EQUAL(*expected_dbl, *actual->field(2));
    }

    TEST_F(DenseUnionBuilderTest, InferredType) {
        AppendInferred();

        auto expected_i8 = assert_array_from_json(int8(), "[33, 10, -10]");
        auto expected_str = assert_array_from_json(utf8(), R"(["abc", "", "def"])");
        auto expected_dbl = assert_array_from_json(float64(), "[1.0, -1.0, 0.5]");

        auto expected_offsets = assert_array_from_json(int32(), "[0, 1, 0, 1, 2, 2, 0, 1, 2]");

        ASSERT_OK_AND_ASSIGN(auto expected,
                             DenseUnionArray::create(*expected_types, *expected_offsets,
                                                     {expected_i8, expected_str, expected_dbl},
                                                     {"i8", "str", "dbl"}, {I8, STR, DBL}));

        ASSERT_EQ(expected->type()->to_string(), actual->type()->to_string());
        ASSERT_ARRAYS_EQUAL(*expected, *actual);
    }

    TEST_F(DenseUnionBuilderTest, ListOfInferredType) {
        std::shared_ptr<ListArray> actual;
        AppendListOfInferred(&actual);

        auto expected_type = list(
                dense_union({field("i8", int8()), field("str", utf8()), field("dbl", float64())},
                            {I8, STR, DBL}));
        ASSERT_EQ(expected_type->to_string(), actual->type()->to_string());
    }

    TEST_F(SparseUnionBuilderTest, Basics) {
        union_builder.reset(new SparseUnionBuilder(
                default_memory_pool(), {i8_builder, str_builder, dbl_builder},
                sparse_union({field("i8", int8()), field("str", utf8()), field("dbl", float64())},
                             {I8, STR, DBL})));

        AppendBasics();

        auto expected_i8 =
                assert_array_from_json(int8(), "[33, null, null, null, null, 10, null, -10, null]");
        auto expected_str =
                assert_array_from_json(utf8(), R"([null, "abc", null, null, "",  null, "def", null, null])");
        auto expected_dbl =
                assert_array_from_json(float64(), "[null, null, 1.0, -1.0, null, null, null, null, 0.5]");

        ASSERT_OK_AND_ASSIGN(
                auto expected,
                SparseUnionArray::create(*expected_types, {expected_i8, expected_str, expected_dbl},
                                         {"i8", "str", "dbl"}, {I8, STR, DBL}));

        ASSERT_EQ(expected->type()->to_string(), actual->type()->to_string());
        ASSERT_ARRAYS_EQUAL(*expected, *actual);
    }

    TEST_F(SparseUnionBuilderTest, NullsAndEmptyValues) {
        union_builder.reset(new SparseUnionBuilder(
                default_memory_pool(), {i8_builder, str_builder, dbl_builder},
                sparse_union({field("i8", int8()), field("str", utf8()), field("dbl", float64())},
                             {I8, STR, DBL})));
        AppendNullsAndEmptyValues();

        // "abc", null, 0, 42, null, null, 0, 0
        // (note that getting 0 for empty values is implementation-defined)
        auto expected_i8 = assert_array_from_json(int8(), "[0, null, 0, 42, null, null, 0, 0]");
        auto expected_str = assert_array_from_json(utf8(), R"(["abc", "", "", "", "", "", "", ""])");
        auto expected_dbl = assert_array_from_json(float64(), "[0, 0, 0, 0, 0, 0, 0, 0]");

        ASSERT_OK_AND_ASSIGN(
                auto expected,
                SparseUnionArray::create(*expected_types, {expected_i8, expected_str, expected_dbl},
                                         {"i8", "str", "dbl"}, {I8, STR, DBL}));

        ASSERT_EQ(expected->type()->to_string(), actual->type()->to_string());
        ASSERT_ARRAYS_EQUAL(*expected, *actual);
        // Physical arrays must be as expected
        ASSERT_ARRAYS_EQUAL(*expected_i8, *actual->field(0));
        ASSERT_ARRAYS_EQUAL(*expected_str, *actual->field(1));
        ASSERT_ARRAYS_EQUAL(*expected_dbl, *actual->field(2));
    }

    TEST_F(SparseUnionBuilderTest, InferredType) {
        AppendInferred();

        auto expected_i8 =
                assert_array_from_json(int8(), "[33, 10, null, null, null, -10, null, null, null]");
        auto expected_str =
                assert_array_from_json(utf8(), R"([null, null, "abc", "", "def",  null, null, null, null])");
        auto expected_dbl =
                assert_array_from_json(float64(), "[null, null, null, null, null, null, 1.0, -1.0, 0.5]");

        ASSERT_OK_AND_ASSIGN(
                auto expected,
                SparseUnionArray::create(*expected_types, {expected_i8, expected_str, expected_dbl},
                                         {"i8", "str", "dbl"}, {I8, STR, DBL}));

        ASSERT_EQ(expected->type()->to_string(), actual->type()->to_string());
        ASSERT_ARRAYS_EQUAL(*expected, *actual);
    }

    TEST_F(SparseUnionBuilderTest, StructWithUnion) {
        auto union_builder = std::make_shared<SparseUnionBuilder>(default_memory_pool());
        StructBuilder builder(STRUCT({field("u", union_builder->type())}),
                              default_memory_pool(), {union_builder});
        ASSERT_EQ(union_builder->AppendChild(std::make_shared<Int32Builder>(), "i"), 0);
        ASSERT_TRUE(builder.type()->equals(
                STRUCT({field("u", sparse_union({field("i", int32())}, {0}))})));
    }

}  // namespace nebula
